import numpy as np


class SimpleLinearRegression:

    def __init__(self):
        self.a_ = None
        self.b_ = None
        self.x_avg_ = None
        self.y_avg_ = None

    def fit(self, x, y):
        assert x is not None and y is not None, \
            'buneng wei kong'
        assert x.shape[1] == y.shape[1], \
            'wei du xiang tong '
        self.x_avg_ = np.mean(x)
        self.y_avg_ = np.mean(y)
        self.a_ = (x - self.x_avg_).dot(y - self.y_avg_) / (x - self.x_avg_).dot((x - self.x_avg_))
        self.b_ = self.y_avg - self.a_ * self.x_avg_
        return self

    def predict(self, x):
        return np.array([self.a_ * i + self.b_ for i in x])

    def __repr__(self):
        return "SimpleLinearRegression"
